use core::{alloc::GlobalAlloc, marker::PhantomData};

use uefi_ffi::services::boot::memory::EfiMemoryType;

use self::pool::Pool;

pub mod heap;
pub mod page;
pub mod pool;

pub trait MemoryType {
    const MEM_TYPE: EfiMemoryType;
}

macro_rules! def_mem_type {
    ($($MEM:tt = $VALUE:path);* $(;)?) => {
        $(
           pub struct $MEM;
            impl MemoryType for $MEM {
                const MEM_TYPE: EfiMemoryType = $VALUE;
            }
        )*
    };
}

def_mem_type! {
    LoaderCode = EfiMemoryType::EfiLoaderCode;
    LoaderData = EfiMemoryType::EfiLoaderData;
    BootServicesCode = EfiMemoryType::EfiBootServicesCode;
    BootServicesData = EfiMemoryType::EfiBootServicesData;
    RuntimeServicesCode = EfiMemoryType::EfiRuntimeServicesCode;
}

pub struct SystemAllocator {
    phantom: PhantomData<usize>,
}

#[global_allocator]
static GLOBAL_ALLOCATOR: SystemAllocator = SystemAllocator {
    phantom: PhantomData,
};

unsafe impl GlobalAlloc for SystemAllocator {
    unsafe fn alloc(&self, layout: core::alloc::Layout) -> *mut u8 {
        Pool::<BootServicesData>::alloc(layout.size())
            .unwrap()
            .as_ptr()
    }

    unsafe fn dealloc(&self, ptr: *mut u8, _layout: core::alloc::Layout) {
        if let Some(pool) = Pool::<BootServicesData>::from_raw(ptr) {
            pool.free()
        }
    }
}
